#include "compiler/semantic_analyzer.h"
#include <iostream>

SemanticAnalyzer::SemanticAnalyzer(std::shared_ptr<Catalog> catalog)
    : catalog(catalog) {
}

void SemanticAnalyzer::analyze(const std::unique_ptr<AST::Statement>& stmt) {
    if (dynamic_cast<AST::CreateTableStmt*>(stmt.get())) {
        analyzeCreateTable(*static_cast<AST::CreateTableStmt*>(stmt.get()));
    }
    else if (dynamic_cast<AST::InsertStmt*>(stmt.get())) {
        analyzeInsert(*static_cast<AST::InsertStmt*>(stmt.get()));
    }
    else if (dynamic_cast<AST::SelectStmt*>(stmt.get())) {
        analyzeSelect(*static_cast<AST::SelectStmt*>(stmt.get()));
    }
    else if (dynamic_cast<AST::DeleteStmt*>(stmt.get())) {
        analyzeDelete(*static_cast<AST::DeleteStmt*>(stmt.get()));
    }
}

void SemanticAnalyzer::analyzeCreateTable(const AST::CreateTableStmt& stmt) {
    if (catalog->getTable(stmt.tableName) != nullptr) {
        throw SemanticError("Table '" + stmt.tableName + "' already exists", stmt.line, stmt.column);
    }

    // Check for duplicate column names
    for (size_t i = 0; i < stmt.columns.size(); ++i) {
        for (size_t j = i + 1; j < stmt.columns.size(); ++j) {
            if (stmt.columns[i].first == stmt.columns[j].first) {
                throw SemanticError("Duplicate column name: " + stmt.columns[i].first, stmt.line, stmt.column);
            }
        }
    }
}

void SemanticAnalyzer::analyzeInsert(const AST::InsertStmt& stmt) {
    const TableInfo* tableInfo = catalog->getTable(stmt.tableName);
    if (tableInfo == nullptr) {
        throw SemanticError("Table '" + stmt.tableName + "' does not exist", stmt.line, stmt.column);
    }

    // If columns are specified, check they exist
    if (!stmt.columns.empty()) {
        for (const auto& colName : stmt.columns) {
            bool found = false;
            for (const auto& colInfo : tableInfo->columns) {
                if (colInfo.name == colName) {
                    found = true;
                    break;
                }
            }
            if (!found) {
                throw SemanticError("Column '" + colName + "' does not exist in table '" + stmt.tableName + "'", stmt.line, stmt.column);
            }
        }
    }

    // Check value count matches column count
    size_t expectedValues = stmt.columns.empty() ? tableInfo->columns.size() : stmt.columns.size();
    if (stmt.values.size() != expectedValues) {
        throw SemanticError("Expected " + std::to_string(expectedValues) + " values, got " +
            std::to_string(stmt.values.size()), stmt.line, stmt.column);
    }

    // Check value types match column types
    for (size_t i = 0; i < stmt.values.size(); ++i) {
        // Get column index (considering if columns are specified in the insert statement)
        size_t colIndex = i;
        if (!stmt.columns.empty()) {
            // Find column index by name
            const std::string& colName = stmt.columns[i];
            bool found = false;
            for (size_t j = 0; j < tableInfo->columns.size(); ++j) {
                if (tableInfo->columns[j].name == colName) {
                    colIndex = j;
                    found = true;
                    break;
                }
            }
            if (!found) {
                // This should have been caught earlier
                throw SemanticError("Column '" + colName + "' does not exist", 0, 0);
            }
        }

        // Get column type
        const std::string& colType = tableInfo->columns[colIndex].type;

        // Check if value is a Literal
        AST::Literal* literal = dynamic_cast<AST::Literal*>(stmt.values[i].get());
        if (!literal) {
            continue;
        }

        // Validate type match
        if (colType == "INT") {
            if (literal->type != TokenType::CONST_INT) {
                throw SemanticError("Type mismatch: column '" + tableInfo->columns[colIndex].name + "' expects INT", 0, 0);
            }
        } else if (colType == "VARCHAR") {
            if (literal->type != TokenType::CONST_STRING) {
                throw SemanticError("Type mismatch: column '" + tableInfo->columns[colIndex].name + "' expects VARCHAR", 0, 0);
            }
        }
        // Add more type checks as needed
    }
}

void SemanticAnalyzer::analyzeSelect(const AST::SelectStmt& stmt) {
    const TableInfo* tableInfo = catalog->getTable(stmt.tableName);
    if (tableInfo == nullptr) {
        throw SemanticError("Table '" + stmt.tableName + "' does not exist", stmt.line, stmt.column);
    }

    // Check if all columns exist
    if (stmt.columns.size() == 1 && stmt.columns[0] == "*") {
        // SELECT * is valid
    }
    else {
        for (const auto& colName : stmt.columns) {
            bool found = false;
            for (const auto& colInfo : tableInfo->columns) {
                if (colInfo.name == colName) {
                    found = true;
                    break;
                }
            }
            if (!found) {
                throw SemanticError("Column '" + colName + "' does not exist in table '" + stmt.tableName + "'", stmt.line, stmt.column);
            }
        }
    }

    // Analyze WHERE clause if present
    if (stmt.whereClause) {
        analyzeExpression(*stmt.whereClause, tableInfo);
    }
}

void SemanticAnalyzer::analyzeDelete(const AST::DeleteStmt& stmt) {
    const TableInfo* tableInfo = catalog->getTable(stmt.tableName);
    if (tableInfo == nullptr) {
        throw SemanticError("Table '" + stmt.tableName + "' does not exist", stmt.line, stmt.column);
    }

    // Analyze WHERE clause if present
    if (stmt.whereClause) {
        analyzeExpression(*stmt.whereClause, tableInfo);
    }
}

void SemanticAnalyzer::analyzeExpression(const AST::Expression& expr, const TableInfo* tableInfo) {
    if (const auto* ident = dynamic_cast<const AST::Identifier*>(&expr)) {
        // Check if identifier is a valid column name
        bool found = false;
        for (const auto& colInfo : tableInfo->columns) {
            if (colInfo.name == ident->name) {
                found = true;
                break;
            }
        }
        if (!found) {
            throw SemanticError("Column '" + ident->name + "' does not exist in table '" + tableInfo->name + "'", ident->line, ident->column);
        }
    }
    else if (const auto* binExpr = dynamic_cast<const AST::BinaryExpression*>(&expr)) {
        analyzeExpression(*binExpr->left, tableInfo);
        analyzeExpression(*binExpr->right, tableInfo);
    }
    // Literals don't need validation
}